package fileIO;

import java.awt.Image;
import java.awt.image.PixelGrabber;
import java.io.File;
import java.io.IOException;
import java.io.FileNotFoundException;
import java.io.FileWriter;
import java.io.FilenameFilter;
import java.util.ArrayList;
import java.util.List;
import javax.imageio.ImageIO;

/**
 *
 * @author Peter Schmidt
 */
public class ImageToFANN {

    // This property will vary between Windows/Unix/Linux/OSX, i.e. it is not always "\n"
    private static final String lineSeparator = System.getProperty("line.separator");

    /**
     * Configures some parameters, essentially the program defaults,
     * and builds the FANN input files needed
     */
    public static void main(String[] args)
            throws FileNotFoundException, IOException, InterruptedException {

        String trainingDataLocation = "training_data";
        String testDataLocation = "test_data";
        Double proportion = 0.05;
        OutputOrder outOrder = OutputOrder.MixedOrder;
        List<String> classes = new ArrayList<String>();

        for (String s : args) {
            if (s.equals("-h") || s.equals("-help") || s.equals("--help")) {
                printHelp();
                System.exit(0);
            } else if (s.startsWith("-train=")) {
                s = s.replace("-train=", "");
                trainingDataLocation = s;
            } else if (s.startsWith("-test=")) {
                s = s.replace("-test=", "");
                testDataLocation = s;
            } else if (s.startsWith("-p=") || s.startsWith("-proportion=")) {
                s = s.replace("-proportion=", "");
                s = s.replace("-p=", "");
                proportion = Double.parseDouble(s);
            } else if (s.startsWith("-o=") || s.startsWith("-order=")) {
                s = s.replace("-order=", "");
                s = s.replace("-o=", "");
                s = s.toLowerCase();
                if (s.equals("posneg")) {
                    outOrder = OutputOrder.PosNegOrder;
                } else if (s.equals("negpos")) {
                    outOrder = OutputOrder.NegPosOrder;
                } else if (s.equals("mixed")) {
                    outOrder = OutputOrder.MixedOrder;
                } else if (s.equals("uniform")) {
                    outOrder = OutputOrder.UniformOrder;
                } else {
                    System.out.println("Order must be one of posneg, negpos, mixed or uniform");
                    System.exit(1);
                }
            } else if (s.startsWith("-c=") || s.startsWith("-charset=")) {
                s = s.replace("-charset=", "");
                s = s.replace("-c=", "");

                if (s.contains("a")) {
                    for (char c = 'a'; c <= 'z'; c++) {
                        classes.add(String.valueOf(c));
                    }
                }
                if (s.contains("A")) {
                    for (char c = 'A'; c <= 'Z'; c++) {
                        classes.add(String.valueOf(c));
                    }
                }
                if (s.contains("0")) {
                    for (char c = '0'; c <= '9'; c++) {
                        classes.add(String.valueOf(c));
                    }
                }
            }
        }

        if (classes.isEmpty()) {
            for (char c = 'a'; c <= 'z'; c++) {
                classes.add(String.valueOf(c));
            }
            for (char c = 'A'; c <= 'Z'; c++) {
                classes.add(String.valueOf(c));
            }
            for (char c = '0'; c <= '9'; c++) {
                classes.add(String.valueOf(c));
            }
        }

        List<String> namedClasses = new ArrayList<String>();
        for (String s : classes) {
            namedClasses.add(folderNameBuilder(s));
        }
        for (String classification : namedClasses) {
            final String fileExt = ".gif",
                    fannExt = ".data";
            String PositiveInDir = String.format(
                    "%s/%s", trainingDataLocation, classification),
                    outFilePath = String.format(
                    "%s/%s%s", trainingDataLocation, classification, fannExt);

            System.out.println(String.format("Building file: %s", outFilePath));
            buildFANNInputFile(trainingDataLocation,
                    PositiveInDir, fileExt, outFilePath, outOrder, proportion);
            System.out.println(String.format("Successfully built file: %s", outFilePath));


            PositiveInDir = String.format(
                    "%s/%s", testDataLocation, classification);
            outFilePath = String.format(
                    "%s/%s%s", testDataLocation, classification, fannExt);
            System.out.println(String.format("Building file: %s", outFilePath));
            buildFANNInputFile(testDataLocation,
                    PositiveInDir, fileExt, outFilePath, outOrder, proportion);
            System.out.println(String.format("Successfully built file: %s", outFilePath));

        }

    }

    /**
     * Work around Windows/Mac SVN issue with file formats.
     * Converts any lower case character to lower case with an underscore 
     * suffix. E.g. "A" remains unchanged as it is upper case,
     * while "a" becomes "a_" as it is lower case.
     */
    static String folderNameBuilder(String name) {
        char[] c = name.toCharArray();
        StringBuilder s = new StringBuilder();
        for (int i = 0; i < c.length; i++) {
            if (String.valueOf(c[i]).matches("[a-z]")) {
                s.append(c[i]);
                s.append('_');
            } else {
                s.append(c[i]);
            }
        }
        return s.toString();
    }

    /**
     * Takes the required input directory locations,
     * searches them for positive and negative training examples.
     * converts the pixel information to black ('0') or white ('1'),
     * and generates the output file (a.k.a. FANN input file) at the
     * specified output file location.
     *
     * @param trainingDataLocation The location containing directories of 
     * positive and negative training examples, grouped into directories by
     * classification
     * @param positiveInDir The directory within training data that contains
     * the positive example images
     * @param inFileExt The extension of the input image files to look for.
     * Must be a literal, e.g. ".bmp", ".gif", ".jpg" or ".png"
     * are acceptable, but not "[.bmp|.gif]"
     * @param outFilePath The location to store the output FANN input file.
     * @throws FileNotFoundException
     * @throws IOException
     * @throws InterruptedException 
     */
    static void buildFANNInputFile(
            final String trainingDataLocation,
            final String positiveInDir,
            final String inFileExt,
            final String outFilePath,
            final OutputOrder outputOrder,
            final Double proportion)
            throws FileNotFoundException, IOException, InterruptedException {

        // Sanitise inputs
        if (positiveInDir == null || positiveInDir.isEmpty()) {
            throw new IllegalArgumentException("Positives folder must not be null or empty");
        }
        if (positiveInDir.endsWith("/") || positiveInDir.endsWith("\\")) {
            throw new IllegalArgumentException("Positives folder must not end with a slash / or \\");
        }

        // Override file name filter to only accept provided suffix/extension
        FilenameFilter filterWhiteListExtension = new FilenameFilter() {

            @Override
            public boolean accept(File dir, String name) {
                return name.endsWith(inFileExt);
            }
        };

        // Initialise
        final int outNeuronSize = 1;
        int exampleCount = 0;

        // Set up the positve examples
        List<File> positiveExamples = new ArrayList<File>();
        File dir = new File(positiveInDir);
        String[] posChildren = dir.list(filterWhiteListExtension);
        for (String child : posChildren) {
            final File inFile = new File(String.format("%s/%s",
                    positiveInDir, child));
            positiveExamples.add(inFile);
        }

        // Set up the negative dirs for the example
        List<File> negativeInDirs = new ArrayList<File>();
        File allDirs = new File(trainingDataLocation);
        final String[] ss = allDirs.list();
        for (String s : ss) {
            final String t = String.format("%s/%s", trainingDataLocation, s);
            if (!t.equals(positiveInDir)) {
                File file = new File(t);
                if (file.isDirectory() && !file.isHidden()) {
                    negativeInDirs.add(file);
                }
            }
        }

        // Randomly sample with replacement the given proportion, rounding down
        List<File> negativeExamples = new ArrayList<File>();
        for (File negDir : negativeInDirs) {
            String[] children = negDir.list(filterWhiteListExtension);
            int sampleSize = (int) ((children.length) * proportion);
            if (sampleSize == 0) {
                sampleSize = 1;
            }
            exampleCount += sampleSize;
            for (int i = 0; i < sampleSize; i++) {
                final Double rand = Math.random() * children.length;
                final int rnd = rand.intValue();
                final File inFile = new File(String.format("%s/%s",
                        negDir, children[rnd]));
                negativeExamples.add(inFile);
            }
        }

        // Clean up existing file if present
        File old = new File(outFilePath);
        if (old.exists()) {
            old.delete();
        }

        FileWriter out = null;
        try {
            // Open new filewriter to append to
            out = new FileWriter(outFilePath, true);

            // Assume first children identical, determine and write additional metadata
            File inFile = positiveExamples.get(0);
            final Image img = ImageIO.read(inFile);
            final int h = img.getHeight(null);
            final int w = img.getWidth(null);

            exampleCount += positiveExamples.size();    // Add +ve examples to count

            out.write(String.valueOf(exampleCount));
            out.write(" ");
            out.write(String.valueOf(h * w));
            out.write(" ");
            out.write(String.valueOf(outNeuronSize));
            out.write(lineSeparator);

            switch (outputOrder) {
                case PosNegOrder:
                    writePosNegOrder(out, positiveExamples, negativeExamples);
                    break;
                case NegPosOrder:
                    writeNegPosOrder(out, positiveExamples, negativeExamples);
                    break;
                case MixedOrder:
                    writeMixedOrder(out, positiveExamples, negativeExamples);
                    break;
                case UniformOrder:
                    writeUniformOrder(out, positiveExamples, negativeExamples);
                    break;
                default:
                    throw new IllegalArgumentException("Output order must be specified");
            }
        } finally {
            if (out != null) {
                out.close();
            }
        }
    }

    /**
     * Uses the passed FileWrtier, generates the file output
     * as all positive examples, then all negative examples
     */
    static void writePosNegOrder(
            FileWriter out,
            List<File> positiveExamples,
            List<File> negativeExamples)
            throws IOException, InterruptedException {

        // Write positives
        for (File inFile : positiveExamples) {
            final String s = buildOutputFromImage(inFile);
            out.write(s);
            out.write(lineSeparator);
            out.write("1");     // finish positive example output neuron
            out.write(lineSeparator);
        }

        // Write negatives
        for (File inFile : negativeExamples) {
            final String t = buildOutputFromImage(inFile);
            out.write(t);
            out.write(lineSeparator);
            out.write("0");     // finish negative example output neuron
            out.write(lineSeparator);
        }
    }

    /**
     * Uses the passed FileWrtier, generates the file output
     * as all negative examples, then all positive examples
     */
    static void writeNegPosOrder(
            FileWriter out,
            List<File> positiveExamples,
            List<File> negativeExamples)
            throws IOException, InterruptedException {

        // Write negatives
        for (File inFile : negativeExamples) {
            final String t = buildOutputFromImage(inFile);
            out.write(t);
            out.write(lineSeparator);
            out.write("0");     // finish negative example output neuron
            out.write(lineSeparator);
        }

        // Write positives
        for (File inFile : positiveExamples) {
            final String s = buildOutputFromImage(inFile);
            out.write(s);
            out.write(lineSeparator);
            out.write("1");     // finish positive example output neuron
            out.write(lineSeparator);
        }

    }

    /**
     * Uses the passed FileWrtier, generates the file output
     * as writing alternating positive and negative examples,
     * with excess negative or positive examples grouped in the middle.
     */
    static void writeMixedOrder(
            FileWriter out,
            List<File> positiveExamples,
            List<File> negativeExamples)
            throws IOException, InterruptedException {

        // Write first half of positives and negatives
        final int firstHalf = Math.min(negativeExamples.size(), positiveExamples.size()) / 2;
        for (int i = 0; i < firstHalf; i++) {
            File inFile = positiveExamples.get(0);
            out.write(buildOutputFromImage(inFile));
            out.write(lineSeparator);
            out.write("1");
            out.write(lineSeparator);
            positiveExamples.remove(inFile);

            inFile = negativeExamples.get(0);
            out.write(buildOutputFromImage(inFile));
            out.write(lineSeparator);
            out.write("0");
            out.write(lineSeparator);
            negativeExamples.remove(inFile);
        }

        // Equalise number of negatives and positives
        while (negativeExamples.size() > positiveExamples.size()) {
            File inFile = negativeExamples.get(0);
            out.write(buildOutputFromImage(inFile));
            out.write(lineSeparator);
            out.write("0");
            out.write(lineSeparator);
            negativeExamples.remove(inFile);
        }
        while (negativeExamples.size() < positiveExamples.size()) {
            File inFile = positiveExamples.get(0);
            out.write(buildOutputFromImage(inFile));
            out.write(lineSeparator);
            out.write("1");
            out.write(lineSeparator);
            positiveExamples.remove(inFile);
        }

        // Write remaining positives and negatives
        while (positiveExamples.size() > 0) {
            File inFile = positiveExamples.get(0);
            out.write(buildOutputFromImage(inFile));
            out.write(lineSeparator);
            out.write("1");
            out.write(lineSeparator);
            positiveExamples.remove(inFile);

            inFile = negativeExamples.get(0);
            out.write(buildOutputFromImage(inFile));
            out.write(lineSeparator);
            out.write("0");
            out.write(lineSeparator);
            negativeExamples.remove(inFile);
        }

    }

    /**
     * Uses the passed FileWrtier, generates the file output
     * as writing alternating positive and negative examples such
     * that the distribution is uniform, e.g. if 4 positives (+) and 
     * 2 negatives (-), it would write + + - + + -
     */
    static void writeUniformOrder(FileWriter out,
            List<File> positiveExamples,
            List<File> negativeExamples)
            throws IOException, InterruptedException {

        // Determine proportion of - to +
        final Double proportions =
                (double) negativeExamples.size()
                / (double) positiveExamples.size();

        // Initialise and run first + since it must exist for a
        // reasonable classification
        int negCount = 0,
                posCount = 1;

        File inFile = positiveExamples.get(0);
        positiveExamples.remove(inFile);

        out.write(buildOutputFromImage(inFile));
        out.write(lineSeparator);
        out.write("1");
        out.write(lineSeparator);

        // While each set not empty print - or + depending 
        // on how close to the proportion we are

        while (positiveExamples.size() > 0 || negativeExamples.size() > 0) {
            // If less than proportions, need more negatives
            if (((double) negCount / (double) posCount) < proportions) {
                inFile = negativeExamples.get(0);
                negativeExamples.remove(inFile);
                negCount++;

                out.write(buildOutputFromImage(inFile));
                out.write(lineSeparator);
                out.write("0");
                out.write(lineSeparator);
                // Otherwise greater than proportions so need more positives
            } else {
                inFile = positiveExamples.get(0);
                positiveExamples.remove(inFile);
                posCount++;

                out.write(buildOutputFromImage(inFile));
                out.write(lineSeparator);
                out.write("1");
                out.write(lineSeparator);
            }
        }
    }

    /**
     * Builds the general output from the image file given, 
     * as needed for the FANN input 
     */
    static String buildOutputFromImage(File inFile)
            throws IOException, InterruptedException {
        StringBuilder s = new StringBuilder();

        final Image img = ImageIO.read(inFile);

        final int h = img.getHeight(null);
        final int w = img.getWidth(null);

        final PixelGrabber pixelGrabber =
                new PixelGrabber(img, 0, 0, w, h, true);
        pixelGrabber.grabPixels();
        int[] pixels = (int[]) pixelGrabber.getPixels();

        for (int x = 0; x < h; x++) {
            for (int y = 0; y < w; y++) {
                int pixel = x * w + y;

                int red = pixels[pixel] & 0x00ff0000 >> 16,
                        green = pixels[pixel] & 0x0000ff00 >> 8,
                        blue = pixels[pixel] & 0x000000ff;
                int greyscale = red + green + blue;

                // if sum less than half call it black
                char blackwhite;
                if (greyscale < (255 * 3 / 2)) {
                    blackwhite = '0';
                } else {
                    blackwhite = '1';
                }
                s.append(blackwhite);
                s.append(' ');
            }
        }
        return s.toString();
    }

    private static void printHelp() {
        System.out.println("ImageToFANN takes the following arguments:");
        System.out.println("-help                \t show this message");
        System.out.println("-train=<path>        \t training data location, e.g. \"../training_data\"");
        System.out.println("-test=<path>         \t test data location");
        System.out.println("-proportion=<double> \t The proportion of negative examples to include\n"
                + "                     \t (will always include at least one of each negative)");
        System.out.println("-charset <0|a|A>     \t The characters to generate files for, where \n"
                + "                     \t 0 = [0-9], a = [a-z], A = [A-Z] in any order");
        System.out.println("-order=<order>       \t Create output files with a specified order where:\n"
                + "\tposneg = positives then negatives, \n"
                + "\tnegpos = negatives then positives, \n"
                + "\tmixed = print positive then negative alternating, \n"
                + "\t\twith excess positives or negatives \n"
                + "\t\tin the middle of the output file \n"
                + "\tuniform = positives and negatives are \n"
                + "\t\tapproximately uniformly distributed\n");
        System.out.println("Defaults are (in order): training_data, test_data, 0.05, 0aA, mixed\n");
    }
}
